#!/usr/bin/python3
# -*- coding: utf-8 -*-

from collections import Iterable
from concurrent.futures import Executor, ProcessPoolExecutor, Future, wait
from datetime import datetime
from logging import getLogger
from multiprocessing import Pipe
from multiprocessing.connection import Connection
from os import environ, getcwd
from os.path import abspath, join as joinpath, normpath

from dateutil import tz

from .callbacks import CallbackLogger, CALLBACK_VERBOSE, CALLBACK_INFO, CALLBACK_SUCCESS, CALLBACK_TITLE, \
    CALLBACK_WARN, CALLBACK_ERROR, CALLBACK_PROGRESS
from .command import Command, Environment, Resource, ensure_dir
from .datetime import now
from .entities import Job, Project
from .envnames import *
from .reflect import create_instance
from .validators import reject_if, reject_if_null

_LOGGER = getLogger(__name__)
_RUNNER_EXECUTOR: Executor = None
_RUNNER_POOL: list = None
_STARTED_TIME_PREFIX = '#started_time,'
_CALLBACK_PREFIX = '#callback,'


class JobRunner(object):
    """启动一个任务，并记录执行任务时的所有日志。"""
    __slots__ = ('__base_dir', '__project', '__job', '__nodes', '__params',
                 '__submitted_time', '__started_time', '__finished_time', '__exits')

    def __init__(
            self,
            project: Project,
            job: Job,
            nodes: Iterable,
            params: dict = None,
    ):
        """构造 `JobRunner` 对象。"""

        reject_if_null(project)
        reject_if_null(job)

        base_dir = environ.get(IKANDA_BASE_DIR, None)
        if not base_dir:
            base_dir = getcwd()
        self.__base_dir = abspath(base_dir)
        self.__project = project
        self.__job = job
        self.__nodes = tuple(nodes)
        self.__params = dict(params) if params else dict()
        self.__submitted_time = None
        self.__started_time = None
        self.__finished_time = None
        self.__exits = []

    @property
    def base_dir(self) -> str:
        return self.__base_dir

    @property
    def project(self) -> Project:
        return self.__project

    @property
    def job(self) -> Job:
        return self.__job

    @property
    def params(self) -> dict:
        """获取执行任务的参数。

        这些参数的定义来自于 `Job.params`，值来自于界面的输入。
        """
        return self.__params

    @property
    def nodes(self) -> tuple:
        return self.__nodes

    @property
    def submitted_time(self) -> datetime:
        return self.__submitted_time

    @submitted_time.setter
    def submitted_time(
            self,
            value: datetime
    ):
        self.__submitted_time = value

    @property
    def started_time(self) -> datetime:
        return self.__started_time

    @property
    def finished_time(self) -> datetime:
        return self.__finished_time

    @finished_time.setter
    def finished_time(
            self,
            value: datetime
    ):
        self.__finished_time = value

    def recv_all(self):
        """从关联的 `Connection` 中读取数据并更新属性。"""
        item = find_by_runner(self)
        if not item:
            raise RuntimeError('cannot find runner in pool, current progress may not be main progress')

        f, r, callback, c1, c2 = item
        while c1.poll():
            s = c1.recv().strip()
            if not s or not s.startswith('#'):
                # 忽略无效的数据。
                _LOGGER.warning('Illegal msg: %s', repr(s))
                continue

            if s.startswith(_STARTED_TIME_PREFIX):
                s = s[len(_STARTED_TIME_PREFIX):]
                self.__started_time = datetime.fromtimestamp(float(s), tz=tz.UTC)
            elif s.startswith(_CALLBACK_PREFIX):
                s = s[len(_CALLBACK_PREFIX):]
                s_timestamp, s_level, s_msg, *etc = (*s.split(',', 2), '', '')
                timestamp = datetime.fromtimestamp(float(s_timestamp), tz=tz.UTC)
                level = int(s_level)
                msg = s_msg
                callback.log(level, msg, timestamp=timestamp)
            else:
                _LOGGER.warning('Illegal msg: %s', repr(s))

    def wait(
            self,
            timeout: (int, None) = None
    ):
        """等待直到此 `JobRunner` 执行完成。如果未指定超时时间，那么一直等待。

        :param timeout: 超时时间，单位是秒。
        """
        item = find_by_runner(self)
        if not item:
            return

        f, *etc = item

        wait((f,), timeout)

    def atexit(
            self,
            fn: object
    ) -> None:
        """添加一个用于在 `JobRunner` 执行结束后运行的函数。

        :param fn: 执行结束后运行的函数。此函数接收四个参数，分别是此 `JobRunner` 本身，执行返回的结果，执行过程中抛出的异常和回调接口对象。
        """
        if self.__exits is None:
            raise RuntimeError('cannot register on-exit function for finished runner')

        self.__exits.append(fn)

    def onexit(
            self,
            result,
            exception: Exception,
            callback: CallbackLogger
    ) -> None:
        """执行所有注册的 `atexit` 函数"""
        if not self.__exits:
            return

        for exit_ in self.__exits:
            exit_(self, result, exception, callback)

        self.__exits = None

    def __getstate__(self):
        """进程间传递 `JobRunner` 需要序列化和反序列化。"""
        return {
            'base_dir': self.__base_dir,
            'project': self.__project,
            'job': self.__job,
            'params': self.__params,
            'nodes': self.__nodes
        }

    def __setstate__(
            self,
            state
    ):
        """进程间传递 `JobRunner` 需要序列化和反序列化。"""
        self.__base_dir = state.get('base_dir', '')
        self.__project = state.get('project', None)
        self.__job = state.get('job', None)
        self.__params = state.get('params', dict())
        self.__nodes = state.get('nodes', list())


class JobRunnerLog(object):
    """任务产生的日志项。"""
    __slots__ = ('__timestamp', '__level', '__msg',)

    def __init__(
            self,
            timestamp: datetime,
            level: int,
            msg: (str, int, float,)
    ):
        reject_if_null(timestamp)
        reject_if(lambda: level < 0)
        reject_if_null(msg)

        self.__timestamp = timestamp
        self.__level = int(level)
        self.__msg = msg

    @property
    def timestamp(self) -> datetime:
        """获取日志项的时间戳。"""
        return self.__timestamp

    @property
    def level(self) -> int:
        """获取日志项的级别。"""
        return self.__level

    @property
    def msg(self) -> (str, int, float,):
        """获取日志项的消息。可能是文本或者数字。"""
        return self.__msg

    @msg.setter
    def msg(
            self,
            value: (str, int, float,)
    ):
        """设置日志项的消息，可能是文本或者数字。"""
        self.__msg = value

    def __str__(self):
        dp = self.timestamp.strftime('%Y-%m-%d %H:%M:%S %z')
        if self.level == CALLBACK_INFO:
            return dp + ' [INFO   ] ' + self.msg
        elif self.level == CALLBACK_VERBOSE:
            return dp + ' [VERBOSE] ' + self.msg
        elif self.level == CALLBACK_SUCCESS:
            return dp + ' [SUCCESS] ' + self.msg
        elif self.level == CALLBACK_TITLE:
            return dp + ' [TITLE  ] ' + self.msg
        elif self.level == CALLBACK_WARN:
            return dp + ' [WARN   ] ' + self.msg
        elif self.level == CALLBACK_ERROR:
            return dp + ' [ERROR  ] ' + self.msg
        elif self.level == CALLBACK_PROGRESS:
            value = '#' * int(float(self.msg) / 2)
            return dp + ' (' + self.msg + '%) ' + value
        else:
            return '<----- unknown level ----->'


class JobRunnerCallback(CallbackLogger):
    """用于记录 `JobRunner` 执行日志的回调类。"""
    __slots__ = ('__logs', '__conn',)

    def __init__(self):
        super().__init__()

        self.__logs = []
        self.__conn = None

    def log_(
            self,
            level: int,
            msg,
            timestamp: datetime
    ):
        if self.__conn:
            # 发送数据到 `Connection` 。
            self.__conn.send(_CALLBACK_PREFIX + str(now().timestamp()) + ',' + str(level) + ',' + str(msg))
        else:
            # 否则记录到列表。
            if level == CALLBACK_PROGRESS:
                last_log = self.__logs[-1] if self.__logs else None
                if last_log.level == level:
                    last_log.msg = msg
                    return

            self.__logs.append(JobRunnerLog(timestamp=timestamp, level=level, msg=msg))

    @property
    def all_logs(self) -> tuple:
        """获取所有执行中产生的日志记录。"""
        return tuple(self.__logs)

    def attach_conn(
            self,
            conn: Connection
    ) -> None:
        """关联到连接。

        如果此记录日志的回调类被关联到了一个 `Connection`，那么所有日志不会被保存，而是通过此连接发送到另一个进程。
        通常在子进程中关联 `Connection`，因为子进程中记录的日志无法被主进程直接读取，所以需要发送到 `Connection` 中，然后再从主进程读取。

        :param conn: 被关联的连接。
        """
        reject_if_null(conn)

        self.__conn = conn


def _create_command(
        command: dict
) -> (Command, None):
    """创建 `Command` 对象。

    :param command: 用于创建 `Command` 对象的字典，包含了类型名和调用构造方法需要的参数。
    :return: 已创建的 `Command` 对象。
    """
    if not command:
        return None

    class_name = command.get('class', '')
    if not class_name:
        raise ValueError('class name must be specified')

    init_params = dict((key, value) for key, value in command.items() if key != 'class')
    return create_instance(class_name, init_params)


def _setup_environment(
        base_dir: str,
        project: Project,
        params: dict,
        callback: CallbackLogger
) -> Environment:
    """准备所有 `Command` 的运行环境。"""
    wd = normpath(joinpath(base_dir, '.workspace'))
    ensure_dir(wd)
    environment = Environment()
    environment.update({
        JOB_WORKING_DIRECTORY: wd,
        JOB_TIMESTAMP: now(),
        PROJECT_NAME: project.name
    })
    if params:
        environment.update(params)
    return environment


def _setup_resources(
        base_dir: str,
        project: Project,
        params: dict,
        callback: CallbackLogger
) -> list:
    """准备所有 `Command` 运行所需的资源。"""
    return []


def _teardown_resources(
        resources: list,
        callback: CallbackLogger
) -> None:
    """所有的 `Command` 运行结束，释放运行过程中创建的资源。"""
    for resource in resources:
        if isinstance(resource, Resource):
            resource.close(callback)


def _start_runner(
        runner: JobRunner,
        callback: JobRunnerCallback,
        conn: Connection
) -> None:
    """执行一个任务。

    :param runner: 已提交的，表示一次执行过程的对象。
    :param callback: 用于记录日志的回调接口。
    :param conn: 用于向主进程发送数据的链接。
    """
    conn.send(_STARTED_TIME_PREFIX + str(now().timestamp()))

    callback.attach_conn(conn)
    environment = _setup_environment(runner.base_dir, runner.project, runner.params, callback)
    resources = _setup_resources(runner.base_dir, runner.project, runner.params, callback)

    callback.info('working directory: {dir}', dir=environment.getstr(JOB_WORKING_DIRECTORY))

    try:
        for command in runner.job.commands:
            command_obj = _create_command(command)
            callback.title(command_obj.description)
            if command_obj.node_dep:
                for node in runner.nodes:
                    if node:
                        callback.title('target node {name} start', name=repr(node.name))
                        environment_cp = Environment(environment)
                        # 向环境变量中加入节点相关的环境变量。
                        environment_cp[NODE_NAME] = node.name
                        environment_cp[NODE_HOSTNAME] = node.hostname
                        environment_cp[NODE_BUNDLE] = node.bundle
                        for key, value in node.props.items():
                            if value:
                                environment_cp['node.' + key] = value
                        command_obj(callback=callback, environment=environment_cp, resources=resources)
                        callback.success('target node {name} complete', name=repr(node.name))
            else:
                command_obj(callback=callback, environment=environment, resources=resources)
    finally:
        _teardown_resources(resources, callback)


def _dummy_start_runner(
        runner: JobRunner,
        callback: JobRunnerCallback,
        conn: Connection
) -> None:
    """模拟执行一个任务。

    :param runner: 已提交的，表示一次执行过程的对象。
    :param callback: 用于记录日志的回调接口。
    :param conn: 用于向主进程发送数据的链接。
    """
    conn.send(_STARTED_TIME_PREFIX + str(now().timestamp()))

    callback.attach_conn(conn)

    callback.info('base-dir={base_dir}', base_dir=runner.base_dir)
    callback.info('project-name={project_name}', project_name=runner.project.name)
    callback.info('commands=Commands[{commands_len}]', commands_len=len(runner.job.commands))
    callback.info('nodes=[{nodes}]', nodes=','.join(str(node) for node in runner.nodes))
    callback.info('params={params}', params=runner.params)


def init_runner_pool(max_workers: int = None) -> None:
    """初始化任务执行池。

    :param max_workers: 最多可以并行的任务数。
    """
    global _RUNNER_EXECUTOR, _RUNNER_POOL

    max_workers = int(max_workers) if max_workers else 2
    _RUNNER_EXECUTOR = ProcessPoolExecutor(max_workers=max_workers)
    _RUNNER_POOL = []


def list_runners() -> tuple:
    """列出任务执行池的所有任务（包括执行中和排队等待的）

    :return: 任务执行池的所有任务（包括执行中和排队等待的）
    """
    global _RUNNER_POOL

    return tuple(item for item in _RUNNER_POOL)


def find_by_runner(
        runner: JobRunner
) -> tuple:
    """找到关联到指定 `JobRunner` 的对象，包括 `Future`, `JobRunner`, `Callback`, `Connection1`, `Connection2`。

    :param runner: 表示一个任务执行过程。
    """
    reject_if_null(runner)

    global _RUNNER_POOL

    return next((item for item in _RUNNER_POOL if item[1] == runner), None)


def dispose_runner_pool() -> None:
    """关闭任务执行池。"""
    global _RUNNER_EXECUTOR, _RUNNER_POOL

    _RUNNER_EXECUTOR.shutdown(wait=True)
    _RUNNER_POOL.clear()


def submit_runner(
        runner: JobRunner
) -> None:
    """提交一个待执行的任务。

    :param runner: 表示一个任务执行过程。
    """
    reject_if_null(runner)

    global _RUNNER_EXECUTOR, _RUNNER_POOL

    def _cleanup() -> None:
        """执行任务执行过程完成后的清理工作。"""
        runner.recv_all()
        c1.close()
        c2.close()
        runner.finished_time = now()

        _RUNNER_POOL.remove((f, runner, callback, c1, c2,))

        runner.onexit(f.result(), f.exception(), callback)

    callback = JobRunnerCallback()
    c1, c2 = Pipe(False)
    f: Future = _RUNNER_EXECUTOR.submit(_start_runner, runner=runner, callback=callback, conn=c2)
    _RUNNER_POOL.append((f, runner, callback, c1, c2,))
    runner.submitted_time = now()
    f.add_done_callback(lambda f_: _cleanup())
